import xarray as xr
import wrf
import netCDF4 as nc
import numpy as np
import os 
import pandas as pd
from scipy.interpolate import griddata
import metpy
from metpy.units import units 
from scipy.integrate import trapz



import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy.mpl.ticker as cticker  # 导入刻度格式化器
import matplotlib.colors as mcolors






#########################  filetree to storage pics ##############
path=os.path.dirname(__file__)
def makefiletree(path,casename):
    path+=f"{cat}"+casename
    if not os.path.exists(path):
        os.mkdir(path)
    model=[f'{cat}Horizontal',f'{cat}Vertical']
    for m in model:
        if not os.path.exists(path+m):
            os.mkdir(path+m)   
        if m==f'{cat}Horizontal':
            fig=[f'{cat}D',f'{cat}qflux',f'{cat}leading_stream',f'{cat}theta',f'{cat}thetase',f'{cat}energy']
            for ff in fig:
                f=path+m+ff
                if not os.path.exists(f):
                    os.mkdir(f)
        elif m==f'{cat}Vertical':
            fig=[f'{cat}PV',f'{cat}RH',f'{cat}theta_se']
            for ff in fig:
                f=path+m+ff
                if not os.path.exists(f):
                    os.mkdir(f)                       




#####################Class settings ##############################
class metplot:
    def __init__(self):
        fig, ax = plt.subplots(subplot_kw={'projection': ccrs.PlateCarree()},figsize=(15, 10),dpi=400)
        self.fig=fig
        self.ax=ax
        #rainsetting
        self.rain_levels=[0.1,10.0,25.0,50.0,100.0,250.0,500.0]#雨量等级
        raindict=['#A6F28F','#3DBA3D','#61BBFF','#0000FF','#FA00FA','#800040']#颜色列表
        self.rain_cmap=mcolors.ListedColormap(raindict)#产生颜色映射
        self.rain_norm=mcolors.BoundaryNorm(self.rain_levels,self.rain_cmap.N)#生成索引
        #snow setting
        self.snow_levels = [0, 0.1, 2.5, 5, 10, 25, 50] #雪量
        snowcolors = ["#ffffff", "#c5e7b0", "#68b165", "#329ed8", "#286ab3", "#ec46f6"]
        self.snow_cmap = mcolors.ListedColormap(snowcolors)
        self.snow_norm=mcolors.BoundaryNorm(self.snow_levels,self.snow_cmap.N)#生成索引
        
        #底图；lon,lat,title代表底图范围和标题（画图函数可以比底图大，但只显示地图）
    def base(self,lon,lat,title=None):
        #title+map
        self.ax.add_feature(cfeature.COASTLINE) #海岸线
        self.ax.add_feature(cfeature.BORDERS, linestyle=':')   #国界
        china_provinces = cfeature.NaturalEarthFeature(category='cultural', name='admin_1_states_provinces_lines', scale='10m', facecolor='none')
        self.ax.add_feature(china_provinces, edgecolor='black', linewidth=0.5) #省界
        self.ax.set(xlim=(np.min(lon),np.max(lon)))
        self.ax.set(ylim=(np.min(lat),np.max(lat)))    
        self.ax.set_xticks(np.linspace(np.min(lon), np.max(lon), 6))  
        self.ax.xaxis.set_major_formatter(cticker.LongitudeFormatter())  # 格式化为经度格式
        self.ax.set_yticks(np.linspace(np.min(lat), np.max(lat), 4))  
        self.ax.yaxis.set_major_formatter(cticker.LatitudeFormatter())  # 格式化为纬度格式
        self.ax.set_title(title)
    
        
    def contourf(self,lon,lat,ctrf,cmap="bwr",levels=None,norm=None,shrink=0.7):
        contourf=self.ax.contourf(lon,lat,ctrf,cmap=cmap,alpha=1,levels=levels,extend="both",norm=norm)
        self.fig.colorbar(contourf, ax=self.ax, shrink=shrink, orientation='vertical' )
    
    def contour(self,lon,lat,ctr,levels=None,norm=None,colors='black'):
        contour = self.ax.contour(lon, lat , ctr, levels=levels, norm=norm, colors=colors)  # 绘制等高线图
        self.ax.clabel(contour, inline=True, fontsize=10)  # 为等高线添加标签
    
    def barb(self,lon,lat,u,v,barbcolor=None,length=6,regrid_shape=10):
        sizes = dict(emptybarb=0.1, spacing=0.15, height=0.6)  #height 外凸长度；empytbarb0风速大小；spacing间隔；值均是相对于length的比例
        self.ax.barbs(lon, lat, u, v, transform=ccrs.PlateCarree(),
              sizes=sizes, length=length,barb_increments={'half':2,'full':4,'flag':20}, #length长度
              linewidth=1.6, regrid_shape=regrid_shape,barbcolor=barbcolor)
    # WRF点太密，对barb quiver图做下采样    
    def sample_factor(self,lon,lat):
        n=lon.ndim
        match n:
            case 2:
                ny,nx=lon.shape
                ns=np.min([ny,nx])/10
            case 1:
                nx=lon.shape[0]
                ny=lat.shape[0]
                ns=np.min([ny,nx])/15
            case _:
                print('!!!!!! new situation appears!!!!! need to update sampling function !!!!!')
                os._exit(0)
        return int(ns)
        
    def quiver(self,lon,lat,u,v,color='green',ref=0.01,unit=None,scale=5e-5,width=0.003):
        n=self.sample_factor(lon, lat)
        match lon.ndim:
            case 2:        
                quiver=self.ax.quiver(lon[::n,::n],lat[::n,::n], u[::n,::n], v[::n,::n],color=color,scale=scale,scale_units='dots',width=width)
            case 1:
                quiver=self.ax.quiver(lon[::n],lat[::n], u[::n,::n], v[::n,::n],color=color,scale=scale,scale_units='dots',width=width)
        self.ax.quiverkey(quiver,                      #传入quiver句柄
          X=0.9, Y = 0.9,       #确定 label 所在位置，都限制在[0,1]之间
          U = ref,                    #参考箭头长度 
          angle = 0,            #参考箭头摆放角度。默认为0，即水平摆放
          label=f'{ref} {unit}',        #箭头的补充：label的内容  + 
          labelpos='S',          #label在参考箭头的哪个方向; S表示南边
          color = color,labelcolor = 'b', #箭头颜色 + label的颜色
          ) 

    def save(self,path,title):
        plt.savefig(f'{path}{cat}{title}',dpi=400, bbox_inches='tight')
        plt.clf()
        
    def glance(self,var):
        os.system("mkdir glance") ; os.chdir('glance')
        if var.ndim==2:
            plt.contourf(var);plt.colorbar()
            plt.savefig("1")
            plt.clf()
        elif var.ndim==3:
            for p in range(var.shape[0]):
                plt.contourf(var[p,:,:]);plt.colorbar()
                plt.savefig(f"{p}")
                plt.clf()    
                
                
class ERA5TY(metplot): #使用ERA5资料绘制台风类
    def __init__(self,era5path,plotlevels,trange=None,latrange=None,lonrange=None):        
        self.ds = xr.open_dataset(era5path, engine="netcdf4")
        # 根据用户提供的参数裁剪数据
        if latrange is not None :
            self.ds = self.ds.sel(latitude=slice(latrange[1], latrange[0]) )
        if lonrange is not None :
            self.ds=self.ds.sel( longitude=slice(lonrange[0], lonrange[1] ) )
            
        self.lon=self.ds.longitude.values
        self.lat=self.ds.latitude.values
        self.P=self.ds.pressure_level.values
        self.time=[str(t)[:13]  for t in self.ds.valid_time.values]
        self.plotlevels=plotlevels
        self.trange=trange
        if trange==None:    
            self.trange=range(len(self.time))

        
    def geth(self,var,t,p):
        return var[t,:,:,:].loc[p,:,:]
            
    def track(self,t,lat,lon):  #lon,lat -上一时刻中心点
        data=self.geth('z',t,1000)
        #aiming to locate typhone by lowest geoheight. data must be 2D.
        radius=1.5
        data_range = data.loc[lat+radius:lat-radius,lon-radius:lon+radius]
        min_value=np.min(data_range)
        min_position = np.where(data_range == min_value) 
        y_idx=min_position[0][0]
        x_idx=min_position[1][0]
        lat_min=data_range.latitude[y_idx].data
        lon_min=data_range.longitude[x_idx].data
        return lat_min,lon_min
            
    
    def dvgH(self):
        # getvar
        z=self.ds.z
        u=self.ds.u
        v=self.ds.v
        d=self.ds.d
        #
        for p in self.plotlevels:
            zz=z.loc[:,p,:,:]
            dd=d.loc[:,p,:,:]
            levelsz=np.linspace(np.min(zz.values),np.max(zz.values),16)
            levelsd=np.linspace(-np.max(abs(dd))/2,np.max(abs(dd))/2,21)
            for t in range(len(self.trange)):
                super().__init__()
                title=f'{p}hPa_{self.time[t]}_Divergence_wind_geopt'
                self.base(self.lon,self.lat,title=title)
                self.contour(self.lon,self.lat,self.geth(z,t,p),levels=levelsz,colors='blue')
                self.contourf(self.lon,self.lat,self.geth(d,t,p),levels=levelsd,cmap="bwr")
                self.barb(self.lon,self.lat,self.geth(u,t,p),self.geth(v,t,p),barbcolor='black')
                self.save(path+cat+'Freddy'+cat+'Horizontal'+cat+'D',title)
                
    def qfluxH(self):
        z=self.ds.z
        u=self.ds.u
        v=self.ds.v
        q=self.ds.q
        qx=u*q/9.8
        qy=v*q/9.8
        for p in self.plotlevels:
            zz=z.loc[:,p,:,:]
            levelsz=np.linspace(np.min(zz.values),np.max(zz.values),16)
            for t in range(len(self.trange)):
                super().__init__()
                title=f'{p}hPa_{self.time[t]}_vaporflux_geopt'
                self.base(self.lon,self.lat,title=title)
                self.contour(self.lon,self.lat,self.geth(z,t,p),levels=levelsz,colors='blue')
                self.quiver(self.lon,self.lat,self.geth(qx,t,p),self.geth(qy,t,p),unit="kg/m•hPa•s")
                self.save(path+cat+'Freddy'+cat+'Horizontal'+cat+'qflux',title)
                
    def leadingstreamH(self):
        u=self.ds.u.loc[:,500,:,:]
        v=self.ds.v.loc[:,500,:,:]
        z8=self.ds.z.loc[:,850,:,:]
        levels=np.linspace(np.min(z8.values),np.max(z8.values),15 )
        for t in range(len(self.trange)):
            title=f'{self.time[t]}_850geopt_500leadingstream'
            super().__init__()
            self.base(self.lon,self.lat,title)
            self.contour(self.lon,self.lat,z8[t,:,:],levels=levels,colors='blue')
            self.barb(self.lon,self.lat,u[t,:,:],v[t,:,:],barbcolor='black')
            self.save(path+cat+'Freddy'+cat+'Horizontal'+cat+'leading_stream',title)
            
    def thetaH(self):
        #draw
        for p in self.plotlevels:
            for t in range(len(self.trange)):
                # data process
                z=self.geth(self.ds.z, t, p)
                T=self.ds.t[t,:,:,:].loc[p,:,:]
                theta=metpy.calc.potential_temperature(p*units.hPa,T*units.kelvin)  
                title=f'{p}hPa_{self.time[t]}_theta_geopotential'
                # draw
                super().__init__()
                self.base(self.lon,self.lat,title)
                self.contourf(self.lon, self.lat, theta, levels=10)
                self.contour(self.lon,self.lat,z)
                self.save(path+cat+'Freddy'+cat+'Horizontal'+cat+'theta', title)
                
    def pseudo_equivalent_potential_temperature(self,temperature, specific_humidity, pressure):
        R_d = 287.05  # 干空气的气体常数，J/(kg·K)
        c_p = 1005.7  # 干空气的定压比热容，J/(kg·K)
        L_v = 2.5e6   # 水的汽化潜热，J/kg
        p0 = 1000   # 参考气压，1000 hPa
        # 计算混合比 w = q / (1 - q)
        mixing_ratio = specific_humidity / (1 - specific_humidity)
        
        # 计算位温 theta = T * (p0 / p) ** (R_d / c_p)
        theta = temperature * (p0 / pressure) ** (R_d / c_p)
        
        # 假设 T_L ≈ T (近似使用温度代替液态水温度)
        T_L = temperature
        
        # 计算假相当位温 theta_e = theta * exp(L_v * w / (c_p * T_L))
        theta_e = theta * np.exp((L_v * mixing_ratio) / (c_p * T_L))
        return theta_e            
            
    def thetaseH(self):
        #draw
        for p in self.plotlevels:
            for t in range(len(self.trange)):
                # data process
                z=self.geth(self.ds.z, t, p)
                T=self.ds.t[t,:,:,:].loc[p,:,:]
                q=self.ds.q[t,:,:,:].loc[p,:,:]
                thetase=self.pseudo_equivalent_potential_temperature(T,q,p)  
                title=f'{p}hPa_{self.time[t]}_thetase_geopotential'
                # draw
                super().__init__()
                self.base(self.lon,self.lat,title)
                self.contourf(self.lon, self.lat, thetase, levels=10)
                self.contour(self.lon,self.lat,z)
                self.save(path+cat+'Freddy'+cat+'Horizontal'+cat+'thetase', title)             
                
    def energyH(self):
        ##  Energy
        g=9.8
        Cv=718
        l = 2.5e6   # 水的汽化潜热，J/kg
        invp=self.ds.pressure_level[::-1]
        t=self.ds.t
        z=self.ds.z
        u=self.ds.u
        v=self.ds.v
        q=self.ds.q
        I=100*trapz(t,invp,axis=1)/g*Cv #这里需要对pres做逆转处理，否则结果是负数。
        P=trapz(z/g,invp,axis=1)*100
        V=np.sqrt(u**2+v**2)
        K=trapz(V,invp,axis=1)*100/2/g
        S=trapz(q,invp,axis=1)*100/g*l
        lon=self.lon
        lat=self.lat
        for tt in range(len(self.trange)):
            #public thing
            fig, ax4 = plt.subplots(2,2,subplot_kw={'projection': ccrs.PlateCarree()},figsize=(10, 5))
            title1= str(self.time[tt])+"  "
            title2=["Inertial energy","Potential energy","Kinetic energy",'Latent energy']
            for i in range(len(ax4.flat)):
                ax=ax4.flat[i]
                ax.add_feature(cfeature.COASTLINE) #海岸线
                ax.add_feature(cfeature.BORDERS, linestyle=':')
                ax.set(xlim=(min(lon),max(lon)))
                ax.set(ylim=(min(lat),max(lat)))    
                ax.set_xticks(np.linspace(min(lon),max(lon),6))
                ax.xaxis.set_major_formatter(cticker.LongitudeFormatter())  # 格式化为经度格式
                ax.set_yticks(np.linspace(min(lat), max(lat),5))  
                ax.yaxis.set_major_formatter(cticker.LatitudeFormatter())  # 格式化为纬度格式  
                #draw
                title=title1+title2[i]
                ax.set_title(title,fontdict={'size':6})
                if i==0:
                    var=I
                if i==1:
                    var=P
                if i==2:
                    var=K
                if i==3:
                    var=S
                levels=np.linspace(np.min(var),np.max(var),13)
                contourf=ax.contourf(lon,lat,var[tt,:,:],cmap="bwr",alpha=0.7,levels=levels,extend="both")
                fig.colorbar(contourf, ax=ax, shrink=0.7)
            plt.savefig(path+cat+'Freddy'+cat+'Horizontal'+cat+'energy'+cat+title,dpi=400,bbox_inches='tight')
            plt.clf()
    
    



######################Initial settings ##################
era5path=r"C:\Users\lenovo\Desktop\Download\Freddy1.nc"
plotlevels=[150,200,500,700,850,925]  #需要绘制的垂直高度层次；
# plotlevels=[500]
trange=range(24)  # 需要绘制的时次;None是全画；列表形式给出想要的时刻
# trange=[0]
casename="Freddy"  # 创建casename文件放图
lonrange=[20,80]   #默认升序;用于裁剪
latrange=[-40,0]  #默认升序；用于裁剪
cat='\\'  #windows--\\  linxu--/ 
era=ERA5TY(era5path,plotlevels,trange,latrange,lonrange)  #如果trange=None的画，默认所有时次都画
makefiletree(path,casename)


#####################DIY here  ##############################
# era.dvgH()
# era.qfluxH()
# era.leadingstreamH()
# era.thetaH()
era.thetaseH()
# era.energyH()


    
